import logging
import json
import requests
import os

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('mcp.client')

class EmbeddingMCPClient:
    """MCP客户端，用于访问Embedding MCP服务"""
    
    def __init__(self, server_url='http://localhost:8000'):
        """初始化MCP客户端
    
        Args:
            server_url: MCP服务器URL，默认为http://localhost:8000
        """
        self.server_url = server_url
        logger.info(f"MCP客户端初始化，服务器URL: {server_url}")
    
    def process_document(self, file_path=None):
        """处理文档并保存嵌入
        
        Args:
            content: 文本内容
            
        Returns:
            dict: 处理结果，包含status和message
        """
        try:
            url = f"{self.server_url}/api/v1/process_document"
            logger.info(f"发送文档处理请求到: {url}")
            
            # 根据输入类型构建请求
            if file_path is None:
                logger.error(f"未提供文件路径")
                raise ValueError("未提供文件路径")

            # 上传文件
            with open(file_path, 'rb') as f:
                content = f.read()
            # 将文件内容转换为字符串（如果是文本文件）
            # 对于二进制文件，可能需要进行base64编码
            try:
                content_str = content.decode('utf-8')
            except UnicodeDecodeError:
                import base64
                content_str = base64.b64encode(content).decode('utf-8')
                logger.warning(f"文件内容不是纯文本，已进行base64编码")
            
            data = {'content': content_str}
            logger.info(f"上传文件内容大小: {len(content_str)} 字符")
            response = requests.post(url, json=data)
            
            logger.info(f"响应状态码: {response.status_code}")
            logger.info(f"响应内容: {response.text}")
            
            response.raise_for_status()
            
            # 解析响应并确保是有效的字典
            result = response.json()
            if not isinstance(result, dict):
                logger.warning(f"响应不是有效的字典: {result}")
                return {'status': 'success', 'message': '文档处理成功', 'raw_response': result}
            
            return result
        except FileNotFoundError:
            logger.error(f"文件不存在: {file_path}")
            return {
                'status': 'error',
                'message': f'文件不存在: {file_path}'
            }
        except Exception as e:
            logger.error(f"文档处理请求失败: {str(e)}")
            return {
                'status': 'error',
                'message': str(e)
            }
    
    def query(self, query, top_k=3):
        """查询相关文档
        
        Args:
            query: 查询文本
            top_k: 返回结果数量
            
        Returns:
            dict: 查询结果
        """
        try:
            url = f"{self.server_url}/api/v1/query"
            data = {
                'query': query,
                'top_k': top_k
            }
            logger.info(f"发送文档查询请求到: {url}")
            response = requests.post(url, json=data)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"文档查询请求失败: {str(e)}")
            return {
                'status': 'error',
                'message': str(e)
            }
    
    def health_check(self):
        """
        检查服务健康状态
        
        Returns:
            dict: 健康状态
        """
        try:
            url = f"{self.server_url}/health"
            logger.info(f"发送健康检查请求到: {url}")
            response = requests.get(url)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"健康检查请求失败: {str(e)}")
            return {
                'status': 'error',
                'message': str(e)
            }
    
    def process_image_caption(self, image_path, caption):
        """
        处理图片和描述文本的上传
        
        Args:
            image_path: 本地图片文件路径
            caption: 图片描述文本
            
        Returns:
            dict: 处理结果
        """
        try:
            url = f"{self.server_url}/api/v1/process_image_caption"
            logger.info(f"发送图片和描述文本到: {url}")
            
            # 打开图片文件
            with open(image_path, 'rb') as image_file:
                files = {'image': image_file}
                data = {'caption': caption}
                
                response = requests.post(url, files=files, data=data)
                
                logger.info(f"响应状态码: {response.status_code}")
                logger.info(f"响应内容: {response.text}")
                
                response.raise_for_status()
                return response.json()
        except Exception as e:
            logger.error(f"图片和描述文本处理请求失败: {str(e)}")
            return {
                'status': 'error',
                'message': str(e)
            }

# 示例用法
if __name__ == "__main__":
    # 初始化客户端
    client = EmbeddingMCPClient()
    
    # 检查服务健康状态
    health_status = client.health_check()
    logger.info(f"服务健康状态: {health_status}")
    
    # 测试图片和描述文本上传
    # 注意：请将下面的图片路径替换为实际存在的测试图片路径
    image_result = client.process_image_caption(
        image_path="./ide_002.png",
        caption=" IDE使用的AI回答结果"
    )
    logger.info(f"图片和描述文本处理结果: {image_result}")

    # 2. 测试图片描述查询
    # 注意：这里的查询文本应该与之前上传的图片描述相关
    image_query_result = client.query(query="IDE 的使用说明")
    logger.info(f"图片描述查询结果: {image_query_result}")
    

    # 3. 检查是否返回了图片URL
    if image_query_result.get('status') == 'success' and 'results' in image_query_result:
        for i, result in enumerate(image_query_result['results']):
            if 'image_url' in result and result['image_url']:
                logger.info(f"结果 {i+1} 包含图片URL: {result['image_url']}")
                # 尝试访问图片URL
                try:
                    image_response = requests.get(result['image_url'])
                    if image_response.status_code == 200:
                        logger.info(f"成功访问图片URL: {result['image_url']}")
                    else:
                        logger.info(f"访问图片URL失败，状态码: {image_response.status_code}")
                except Exception as e:
                    logger.error(f"访问图片URL时出错: {str(e)}")
            else:
                logger.info(f"结果 {i+1} 不包含图片URL")